%% Step 4

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% AUTHOR: Produced for by JDG/IM based on initial code for cobb-douglas
% production in "Firm Model" directory.
%
% DATE: February 2019.
%
% Purpose: Computes the main algorithm to work out wages, prices, 
% endogenous labour for calibrations with shocks.
%
% Output:  converged expenditure hat, wage hat and prices, 
%         which are calculated via EquilibriumCES.m, which calls up the 
%         following functions:
%           1) PriceLoopCES.m
%           2) TradeShareCES.m
%           3) ExpenditureCES.m
%           4) LabourMarketConditionCES.m
%
% Last Updated: 07/09/2019 by JDG/IM
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

global sigmaT alphaT rhoT lambdaT etaT shockT rho sigma lambda eta ...
    tol maxit Pi_hat_n D_hat_n Xi_hat_mnjf Xi_hat_mnj a_hat_f a_hat ...
    varphi sL_n sPi_n sD_n N J FI_J france oneminus_alpha_vector ...
    labour_share_PWT countrysecD firmsecD_sorted start start_sorted ...
    sum_by_country_dummy ISO alpha_WIOD_j alpha_WIOD_francej sector_algorithm 


%% Algorithm

tol      = 1E-5; %Tolerance allowed until price or wage converged
maxit    = 1E+10; %Maximum iterations to loop over

%% Initialise price hat, wage hat and expenditure hat guesses:
 
% P_hat: Initialise the Price hat equation 

P_hat_mnj0=ones(J*N,N);

% Initialisation of wage shock

w_hat0= ones(1,N);

% Initial guess for the change in expenditures 

X_hat_mnj0=ones(N*J,N);

%% Load in the updated hat data, taking into account the wedges

% Load in the data after completing the convergence the first time, 
% therefore calculating the new hat variables when the wedge is equal to 0. 
% To do counterfactual exercises we start from this point, therefore
% allowing our variables at time t to be equal to the hat variables
% calculated previously. 

%eval([strcat('load MAT/rho',rhoT,'eta',etaT,'lambda',lambdaT,'/Step_4_setup_alpha',num2str(alphaT),...
 %    '_rho',rhoT,'_approx.mat')]);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Homogeneity function
%
% We need to homogeneize (1-alpha)gamma given the structure of the model,
% along with the labor and intermediate shares in the data, for French
% firms. Having done this, we can then run the model as for the
% heterogeneous case.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% X_nkif: Multiplies the expenditure matrix for france with the firm dummy
%         to create a [F x N matrix]

X_nkif = firmsecD_sorted*X_mnj(start(france):start(france+1)-1,:);

% pi_X_nkif: Multiplying firm shares of expenditures (pi_nkj) by the
%            expenditure matrix that has been scaled up by the firm 
%            sector dummy

pi_X_nkif=pi_mnjf.*X_nkif;

% pi_X_nif: This is the sum over country k, should result in a [Fx1] vector

pi_X_nif=sum(pi_X_nkif,2);

%one_minus_pi_l_f: Labor shares for each firm = 1-pi_l_f

one_minus_pi_l_f=1-pi_l_f;

% Multiplying the previous summed equation by this labor share. This is
% done before multiplying by pi_M_f, which will come later.

pi_X_laborshare_nif=pi_X_nif.*one_minus_pi_l_f;

clear pi_X_nkif  X_nkif

pi_X_laborshare_nif_repmat = repmat(pi_X_laborshare_nif, [1 J*N]);     

gamma_pi_X_laborshare_nif = pi_X_laborshare_nif_repmat.*pi_M_f;

% Sum over firms. Should give a J*N matrix, which is numerator of
% homogeneous (1-alpha)*gamma

numerator = firmsecD_sorted'*gamma_pi_X_laborshare_nif;

% Sum exports over destination country per sector for france for numerator
% 32*1 matrix, which must be expanded per destination market

X_fra = X_mnj(start(france):start(france+1)-1,:);
den = sum(X_fra,2);   
denominator = repmat(den,[1 J*N]);

clear pi_X_laborshare_nif_repmat gamma_pi_X_laborshare_nif X_fra 

% Create homogeneous (1-alpha)gamma

one_minus_alp_gam_hom = numerator./denominator;

clear numerator denominator;

% Create homogeneous alpha
pi_X_laborshare_nif=pi_X_nif.*pi_l_f;
numerator = firmsecD_sorted'*pi_X_laborshare_nif;
pi_l_hom = numerator./den;

clear numerator den pi_X_laborshare_nif;
% Create homogeneous gamma given that alpha_hom = alpha_baseline at
% country-sector level

pi_l_hom_check = pi_l(start(france):start(france+1)-1,:);
corr(pi_l_hom,pi_l_hom_check) 
pi_M_hom = one_minus_alp_gam_hom./(1-repmat(pi_l_hom,[1 N*J]));

% Now fill in firm levels given homogeneous values and sectors firms in

%pi_l_f_hom = firmsecD_sorted*pi_l_hom;
pi_l_f_hom = pi_l_f;
temp = firmsecD_sorted*pi_M_hom;
pi_M_f_hom = pi_M_f;
pi_M_f_hom(:,start(idx):start(idx+1)-1) = temp(:,start(idx):start(idx+1)-1);
clear denominator numerator one_minus_alp_gam_hom pi_l_hom pi_M_hom temp;

% Construct homogeneous market shares
N_jf = sum(firmsecD_sorted,1);
N_jf_repmat = repmat(N_jf,FI_J,1);
one_over_N_repmat = firmsecD_sorted ./ N_jf_repmat;
one_over_N = sum(one_over_N_repmat,2);
temp = repmat(one_over_N,1,N);
pi_mnjf_hom = pi_mnjf;
pi_mnjf_hom(:,idx) = temp(:,idx);

% List of non-tradable sectors as given in the calibration.pdf

NT_sectors={'E','F','50','51','52','H','60','61','62','63','64',...
                        'J','70','71t74','L','M','N','O','P'}';

%List of sectors in the WIOD, this is in the order of the WIOD and this is 
%the order we follow    

WIOD_sectors = {'AtB' , 'C' , '15t16' ,'17t18' , '19' , '20' , '21t22' ,...
    '23' , '24' , '25' ,'26', '27t28', '29' , '30t33' , '34t35' ,...
    '36t37', 'E' , 'F' , '50', '51', '52', 'H', '60', '61', '62', '63',...
    '64', '70', '71t74', 'M', 'N', 'O'}';      


for i=1:J
      for j=1:size(NT_sectors,1)
         if isequal(WIOD_sectors{i},NT_sectors{j}) 
            pi_mnjf_hom(start_sorted(i):start_sorted(i+1)-1,1:france-1)=0;  
            pi_mnjf_hom(start_sorted(i):start_sorted(i+1)-1,france+1:N)=0;  
         else
         end
      end
   end


% Function to solve the equilibrium. Note that a change we do here is use
% homogeneous firm-level labour and intermediate shares rather than
% heterogeneous. We are assuming both are homogeneous to begin. Easy to
% create a loop depending on different cases later on.

[P_hat_n,P_hat_nj,P_hat_mnj,w_hat,X_hat_mnj,pi_mnjf1,pi_mnj1,...
    pi_c_mnj1,pi_c_nj1,pi_M1,pi_l1,pi_M_f1,pi_l_f1,PC_hat_n1,...
    final_goods1,intermediates1]= ...
    EquilibriumCES_fun_approx(P_hat_mnj0,w_hat0,pi_mnj,pi_M,pi_M_f_hom,...
    pi_l,pi_l_f_hom,pi_mnjf_hom,pi_c_mnj,X_hat_mnj0,pi_c_nj,PC_n,zeta_mnj,X_mnj);